import paho.mqtt.client as mqtt
import sys
import time
import json
from hashlib import md5

g_config = {}
with open("config.json", "r") as f:
    g_config = json.load(f)

COMMAND_TOPIC = f"shell/{g_config['device_name']}/cmd"
RESPONSE_TOPIC = f"shell/{g_config['device_name']}/rsp"

class MQTTMaster:
    def __init__(self):
        self.client = mqtt.Client()
        self.client.on_connect = self.on_connect
        self.client.on_disconnect = self.on_disconnect
        self.client.on_message = self.on_message
        self.response_received = False
        self.last_response = ""
        self.connected = False
        
    def on_connect(self, client, userdata, flags, rc):
        if rc == 0:
            print("主机端已连接到MQTT Broker")
            self.connected = True
            # 订阅响应主题
            self.client.subscribe(RESPONSE_TOPIC, qos=1)
        else:
            print(f"连接失败，错误代码: {rc}")
            self.connected = False
            
    def on_disconnect(self, client, userdata, rc):
        print("与MQTT Broker断开连接")
        self.connected = False
        
    def on_message(self, client, userdata, msg):
        if msg.topic == RESPONSE_TOPIC:
            self.response_received = True
            try:
                # 解析JSON格式的响应
                json_response = json.loads(msg.payload.decode('utf-8'))
                self.last_response = json_response.get("data", "")
                print(f"{self.last_response}")
            except json.JSONDecodeError:
                # 如果不是JSON格式，显示收到的原始数据
                self.last_response = msg.payload.decode('utf-8')
                print(f"{self.last_response}")
    
    def connect(self):
        try:
            broker_addr = g_config["mqtt_addr"].split(":")[0]
            broker_port = int(g_config["mqtt_addr"].split(":")[1])
            self.client.connect(broker_addr, broker_port, 60)
            self.client.loop_start()
            # 等待连接建立
            time.sleep(1)
            
            # 检查是否连接成功
            if not self.connected:
                print("无法连接到MQTT Broker，请检查网络或Broker设置")
                return False
            return True
        except Exception as e:
            print(f"连接MQTT Broker时出错: {e}")
            return False
    
    def generate_s_data(self, pos=0):
        rand = (int(time.time()) + pos * 3) // 3
        rand_code = g_config["rand_code"]
        s_data = f"{rand_code}{rand}"
        s = md5(s_data.encode('utf-8')).hexdigest()
        return s[:8]

    def verify_s_data(self, s):
        s1 = self.generate_s_data(0) # 当前s
        s2 = self.generate_s_data(-1) # 前一个s
        return s1 == s or s2 == s

    def send_command(self, command):
        if not self.connected:
            print("未连接到MQTT Broker，请先建立连接")
            return False

        self.response_received = False
        self.last_response = ""
        # 将命令包装成JSON格式
        json_command = json.dumps({
            "s": self.generate_s_data(),
            "data": command
        }, ensure_ascii=False)
        
        result = self.client.publish(COMMAND_TOPIC, json_command, qos=1)
        if result.rc != mqtt.MQTT_ERR_SUCCESS:
            print(f"发送命令失败: {result.rc}")
            return False
        return True
        
    def wait_for_response(self, timeout=10):
        start_time = time.time()
        while not self.response_received and (time.time() - start_time) < timeout:
            time.sleep(0.1)
        if not self.response_received:
            print("等待响应超时")
            
    def run_interactive_mode(self):
        print("MQTT Shell 主机端")
        print("输入命令发送到从机，输入 'exit' 退出")
        print("-" * 40)
        
        while True:
            try:
                command = input(">> ")
                if command.lower() == 'exit':
                    break
                if command.strip():
                    if self.send_command(command):
                        self.wait_for_response()
            except KeyboardInterrupt:
                print("收到键盘中断，退出程序")
                break
            except Exception as e:
                print(f"发生错误: {e}")

    def disconnect(self):
        self.client.loop_stop()
        self.client.disconnect()

if __name__ == "__main__":
    master = MQTTMaster()
    if master.connect():
        try:
            master.run_interactive_mode()
        finally:
            master.disconnect()
    else:
        sys.exit(1)